import torch
import torch.nn.functional as F
import numpy as np
from lstm_model import lstm_model

device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model=torch.load('models/宋词.pth')
vocab_size=0

# 获取一个样本
def sample(model, length, top_k=None, word="默认"):
    def predict(model, char, top_k=None, hidden_size=None):

        model.to(device)
        model.eval()
        with torch.no_grad():
            char = np.array([char])  # 转换为array
            char = char.reshape(-1, 1)  # shape转换
            char_encoding = model.onehot_encode(char)  # encoding
            char_encoding = char_encoding.reshape(1, 1, -1)  # (batch_size, seq_len, num_features)
            char_tensor = torch.tensor(char_encoding, dtype=torch.float32)  # 类型转换
            char_tensor = char_tensor.to(device)  # 部署到device上

            out, hidden_size = model(char_tensor, hidden_size)  # 模型预测

            probs = F.softmax(out, dim=1).squeeze()  # torch.Size([1, 83]) --> torch.Size([83])
            # probs = F.softmax(out, dim=1).data # 另一种写法，结果一致

            if top_k is None:
                indices = np.arange(vocab_size)
            else:
                probs, indices = probs.topk(top_k)  # 选取概率最大的前top_k个
                indices = indices.cpu().numpy()

            probs = probs.cpu().numpy()

            char_index = np.random.choice(indices, p=probs / probs.sum())  # 随机选取一个索引
            char = model.int_char[char_index]  # 获取索引对应的字符

        return char, hidden_size
    hidden_size = None  # 初始化
    new_sentence = [char for char in word]  # 初始化
    for i in range(length):
        next_char, hidden_size = predict(model, new_sentence[-1], top_k=top_k, hidden_size=hidden_size)  # 预测下一个字符
        new_sentence.append(next_char)
    result = "".join(new_sentence)
    return result


def 生成藏头诗(top_k,word,sent_len=6):
    result = ""
    for i,sentence in enumerate(word):
        new_text = sample(model, length=sent_len, top_k=top_k,word=sentence)
        print(new_text)
        if i == len(word):
            print(new_text)
            result += new_text


    print(result)

def 生成诗词(top_k,word,len=27):
    new_text = sample(model, length=len, top_k=top_k, word=word)
    result = ""
    for i in range(0, len + 1):
        if i == 27:
            result += new_text[i]
            result += "。"
            break
        if (i + 1) % 7 == 0:
            result += new_text[i]
            result += ",\n"
        else:
            result += new_text[i]

    print(result)

生成藏头诗(top_k=100,word="你无敌了")
print("---------------")
生成诗词(top_k=100,word="气")